/*
 * Copyright 2016-2022 dromara.org.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.dromara.mendmix.gateway.component.ratelimit;

import java.time.Instant;
import java.util.Arrays;
import java.util.List;

import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

/**
 * 
 * @description <br>
 * @author <a href="mailto:vakinge@gmail.com">jiangwei</a>
 * @date 2022年5月21日
 */
public class RedisRatelimiter implements RateLimiter {

	
	private static final String TIMESTAMP_KEY_SUFFIX = "}.timestamp";
	private static final String TOKEN_KEY_SUFFIX = "}.tokens";
	private static final String KEY_PREFIX = "{ratelimit.";

	private static String scriptText = "local tokens_key = KEYS[1]\n" + 
			"local timestamp_key = KEYS[2]\n" + 
			"local rate = tonumber(ARGV[1])\n" + 
			"local capacity = tonumber(ARGV[2])\n" + 
			"local now = tonumber(ARGV[3])\n" + 
			"local requested = 1\n" + 
			"local ttl = math.floor((capacity / rate) * 2)\n" + 
			"local last_tokens = tonumber(redis.call(\"get\", tokens_key)) or capacity\n" + 
			"local last_refreshed = tonumber(redis.call(\"get\", timestamp_key)) or 0\n" + 
			"local delta = math.max(0, now - last_refreshed)\n" + 
			"local filled_tokens = math.min(capacity, last_tokens + (delta * rate))\n" + 
			"local allowed = filled_tokens >= requested\n" + 
			"local new_tokens = filled_tokens\n" + 
			"local allowed_num = 0\n" + 
			"if allowed then\n" + 
			"    new_tokens = filled_tokens - requested\n" + 
			"    allowed_num = 1\n" + 
			"end\n" + 
			"redis.call(\"setex\", tokens_key, ttl, new_tokens)\n" + 
			"redis.call(\"setex\", timestamp_key, ttl, now)\n" + 
			"return allowed_num";
	
	private StringRedisTemplate stringRedisTemplate;
	private DefaultRedisScript<Long> rateLimitRedisScript;
	
	public RedisRatelimiter(StringRedisTemplate stringRedisTemplate) {
		this.stringRedisTemplate = stringRedisTemplate;
		rateLimitRedisScript = new DefaultRedisScript<>();
		rateLimitRedisScript.setScriptText(scriptText);
		rateLimitRedisScript.setResultType(Long.class);
	}

	@Override
	public boolean tryAcquire(String hitKey, double permitsPerSecond) {
		List<String> scriptKeys = getScriptKeys(hitKey);
		int capacity = (int) (permitsPerSecond * 2 + 5);
		long now = Instant.now().getEpochSecond();
		Long allowedNum = stringRedisTemplate.execute(rateLimitRedisScript, scriptKeys, String.valueOf(permitsPerSecond), String.valueOf(capacity), String.valueOf(now));
	    return allowedNum == null || allowedNum > 0;
	}
	
	@Override
	public void resetLimiters() {}

	private static List<String> getScriptKeys(String hitKey) {
		StringBuilder keyBuilder = new StringBuilder(RedisRatelimiter.KEY_PREFIX).append(hitKey);
		int prefixLength = keyBuilder.length();
		String tokenKey = keyBuilder.append(RedisRatelimiter.TOKEN_KEY_SUFFIX).toString();
		keyBuilder.substring(prefixLength);
		String timestampKey = keyBuilder.append(RedisRatelimiter.TIMESTAMP_KEY_SUFFIX).toString();
		return Arrays.asList(tokenKey, timestampKey);
	}
}
